#include <stdio.h>
#include <stdlib.h>
#include <cuda_runtime.h>
#include "debug_utils.h"

#define ARRAY_SIZE 100

// External declaration of the C debug function
extern "C" {
    void debug_print_norm_d(const char* classname, const char* tag, double* array, int size, int rank, int step);
    void debug_print_norm_cuda_d(const char* classname, const char* tag, double* array, int size, int rank, int step);
}

// A simple CUDA kernel that doubles each element (double precision)
__global__ void doubleElementsDouble(double *d_data, int size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size) {
        d_data[idx] = d_data[idx] * 2.0;
    }
}

int main(int argc, char** argv) {
    double* h_data; // Host data
    double* d_data; // Device data
    int rank = 0;   // Use rank 0 for non-MPI tests
    
    // Allocate host memory
    h_data = (double*)malloc(ARRAY_SIZE * sizeof(double));
    if (!h_data) {
        fprintf(stderr, "Host memory allocation failed\n");
        return 1;
    }
    
    // Initialize array with values
    for (int i = 0; i < ARRAY_SIZE; i++) {
        h_data[i] = (double)i / ARRAY_SIZE;
    }
    
    // Print CPU norm before GPU operations
    debug_print_norm_d("CPUdebug_d", "array_values", h_data, ARRAY_SIZE, rank, 0);
    
    // Allocate device memory
    cudaError_t error = cudaMalloc((void**)&d_data, ARRAY_SIZE * sizeof(double));
    if (error != cudaSuccess) {
        fprintf(stderr, "CUDA malloc failed: %s\n", cudaGetErrorString(error));
        free(h_data);
        return 1;
    }
    
    // Copy data from host to device
    error = cudaMemcpy(d_data, h_data, ARRAY_SIZE * sizeof(double), cudaMemcpyHostToDevice);
    if (error != cudaSuccess) {
        fprintf(stderr, "CUDA memcpy H2D failed: %s\n", cudaGetErrorString(error));
        cudaFree(d_data);
        free(h_data);
        return 1;
    }
    
    // Print GPU norm after copying to device
    debug_print_norm_cuda_d("GPUdebug_d", "array_values", d_data, ARRAY_SIZE, rank, 0);
    
    // Launch kernel to double the elements
    int blockSize = 256;
    int numBlocks = (ARRAY_SIZE + blockSize - 1) / blockSize;
    doubleElementsDouble<<<numBlocks, blockSize>>>(d_data, ARRAY_SIZE);
    
    // Check for kernel launch errors
    error = cudaGetLastError();
    if (error != cudaSuccess) {
        fprintf(stderr, "CUDA kernel launch failed: %s\n", cudaGetErrorString(error));
        cudaFree(d_data);
        free(h_data);
        return 1;
    }
    
    // Wait for kernel to finish
    error = cudaDeviceSynchronize();
    if (error != cudaSuccess) {
        fprintf(stderr, "CUDA kernel execution failed: %s\n", cudaGetErrorString(error));
        cudaFree(d_data);
        free(h_data);
        return 1;
    }
    
    // Copy data back from device to host
    error = cudaMemcpy(h_data, d_data, ARRAY_SIZE * sizeof(double), cudaMemcpyDeviceToHost);
    if (error != cudaSuccess) {
        fprintf(stderr, "CUDA memcpy D2H failed: %s\n", cudaGetErrorString(error));
        cudaFree(d_data);
        free(h_data);
        return 1;
    }
    
    // Print CPU norm after GPU operations
    debug_print_norm_d("CPUdebug_d", "array_scaled", h_data, ARRAY_SIZE, rank, 0);
    
    // Print GPU norm after GPU operations
    debug_print_norm_cuda_d("GPUdebug_d", "array_scaled", d_data, ARRAY_SIZE, rank, 0);
    
    // Test with higher precision values to show double advantage
    for (int i = 0; i < ARRAY_SIZE; i++) {
        // Create values that would lose precision in a float
        h_data[i] = 1.0 + (i * 0.00000000001);
    }
    
    // Print CPU norm for high-precision values
    debug_print_norm_d("CPUdebug_d", "high_precision", h_data, ARRAY_SIZE, rank, 1);
    
    // Copy to GPU
    error = cudaMemcpy(d_data, h_data, ARRAY_SIZE * sizeof(double), cudaMemcpyHostToDevice);
    if (error != cudaSuccess) {
        fprintf(stderr, "CUDA memcpy H2D failed: %s\n", cudaGetErrorString(error));
        cudaFree(d_data);
        free(h_data);
        return 1;
    }
    
    // Print GPU norm for high-precision values
    debug_print_norm_cuda_d("GPUdebug_d", "high_precision", d_data, ARRAY_SIZE, rank, 1);
    
    // Modify the array on CPU for a second test
    for (int i = 0; i < ARRAY_SIZE; i++) {
        h_data[i] += 0.5;
    }
    
    // Print norms after CPU modification
    debug_print_norm_d("CPUdebug_d", "array_offset", h_data, ARRAY_SIZE, rank, 2);
    
    // Copy to GPU and back to ensure GPU calculation
    error = cudaMemcpy(d_data, h_data, ARRAY_SIZE * sizeof(double), cudaMemcpyHostToDevice);
    if (error != cudaSuccess) {
        fprintf(stderr, "CUDA memcpy H2D failed: %s\n", cudaGetErrorString(error));
        cudaFree(d_data);
        free(h_data);
        return 1;
    }
    
    // Print GPU norm after second modification
    debug_print_norm_cuda_d("GPUdebug_d", "array_offset", d_data, ARRAY_SIZE, rank, 2);
    
    // Clean up
    cudaFree(d_data);
    free(h_data);
    printf("CUDA double-precision test completed successfully\n");
    printf("Use 'python compare_norm.py' to compare the CPU and GPU norms\n");
    
    return 0;
} 